/**
 * This file is part of the NatCrackerToolbox.
 * 
 * Copyright (C) 2009 - 2011 Peerialism Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * 
 **/
package se.peerialism.natcrackeremu.model.impl;

import static se.peerialism.natcrackeremu.constants.BehaviorConstants.EXPIRATION_TIMER;
import static se.peerialism.natcrackeremu.constants.BehaviorConstants.MAX_PORT;
import static se.peerialism.natcrackeremu.constants.BehaviorConstants.RANDOM_PORT_SEED;
import static se.peerialism.natcrackeremu.constants.BehaviorConstants.TIMER_GC_PERIOD;

import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Random;
import java.util.Set;
import java.util.TimerTask;

import org.apache.log4j.Logger;

import se.peerialism.natcrackeremu.TimerInt;

import com.peerialism.natcracker.common.FilteringPolicy;
import com.peerialism.natcracker.common.PortAllocationPolicy;
import com.peerialism.natcracker.common.PortMappingPolicy;

public class BehavioralNAT {

	class TimerGarbageCollector extends TimerTask {

		@Override
		public void run()
		{
			// System.out.println("GC running");
			Set<Integer> removed = new HashSet<Integer>();
			for (Integer port : timers.keySet())
				if (isExpired(port))
				{
					// Remove from timers
					removed.add(port);
					// Remove from open endpoint
					Endpoint v = pToE.remove(port);
					// Look at reverse mapping
					Map<InetAddress, Map<Integer, Integer>> map = vToDs.get(v);
					for (InetAddress extIp : map.keySet())
					{
						Map<Integer, Integer> iMap = map.get(extIp);
						if (iMap.containsKey(port))
							iMap.remove(port);
					}
					// Remove reverse mapping if empty
					if (map.isEmpty())
						map.remove(v);
				}
			for (Integer port : removed)
				// Remove from timers
				timers.remove(port);
			if (!timers.isEmpty())
				gc.schedule(this, TIMER_GC_PERIOD);
		}
	}
	private static final Logger log = Logger.getLogger(BehavioralNAT.class);
	// NAT Behavior
	private PortMappingPolicy portMappingPolicy;
	private PortAllocationPolicy portAllocationPolicy;
	private final PortAllocationPolicy alternativePortAllocationPolicy;
	private FilteringPolicy filteringPolicy;
	private boolean isExpirationTimerEnabled = true;
	private boolean clashingOverrides = false;
	// NAT State
	private final InetAddress publicAddress;
	/* Open Ports */
	Map<Integer, Endpoint> pToE;
	/* Private to public */
	/*
	 * Private E, Map<Public IP, Map<Public Port, Private
	 * Port>>
	 */
	Map<Endpoint, Map<InetAddress, Map<Integer, Integer>>> vToDs;
	/* Timers */
	Map<Integer, Long> timers;
	// Port Assignment
	private Random portRandom;
	private int portCounter;

	// Timer
	TimerInt gc;

	/**
	 * Creates an instance of this class with the specified
	 * policies.
	 * 
	 * @param portMappingPolicy
	 *            {@link PortMappingPolicy}
	 * @param portAllocationPolicy
	 *            {@link PortAllocationPolicy}
	 * @param filteringPolicy
	 *            {@link FilteringPolicy}
	 * @param publicAddress
	 *            The public {@link InetAddress} address of
	 *            the NAT box
	 * @param timerImplementation
	 *            The {@link TimerInt} implementation
	 *            instance
	 * 
	 */
	public BehavioralNAT(InetAddress publicAddress, PortMappingPolicy portMappingPolicy, PortAllocationPolicy portAllocationPolicy, FilteringPolicy filteringPolicy, TimerInt implementation, PortAllocationPolicy alternativePortAcclocationPolicy) {
		reset();
		this.publicAddress = publicAddress;
		this.portMappingPolicy = portMappingPolicy;
		this.portAllocationPolicy = portAllocationPolicy;
		this.filteringPolicy = filteringPolicy;
		alternativePortAllocationPolicy = alternativePortAcclocationPolicy;
		gc = implementation;
	}

	/**
	 * Creates an instance of this class with default values
	 * for policies. (m=EI,f=EI,a=PP).
	 * 
	 * @param publicAddress
	 *            The public {@link InetAddress} address of
	 *            the NAT box
	 * @param timerImplementation
	 *            The {@link TimerInt} implementation
	 *            instance
	 */
	public BehavioralNAT(InetAddress publicAddress, TimerInt timerImplementation) {
		this(publicAddress, PortMappingPolicy.ENDPOINT_INDEPENDENT,
				PortAllocationPolicy.PORT_PRESERVATION,
				FilteringPolicy.ENDPOINT_INDEPENDENT, timerImplementation,
				PortAllocationPolicy.PORT_CONTIGUITY);
		reset();
	}

	private int getContiguousPort()
	{
		int ret = portCounter++;
		if (portCounter > MAX_PORT)
			ret = portCounter = portRandom.nextInt(MAX_PORT);
		return ret;
	}

	/**
	 * 
	 * @return the external {@link InetAddress} of the
	 *         middlebox
	 */
	public InetAddress getExternalAddress()
	{
		return publicAddress;
	}

	/**
	 * @return the {@link FilteringPolicy}
	 */
	public FilteringPolicy getFilteringPolicy()
	{
		return filteringPolicy;
	}

	/**
	 * @return the {@link PortAllocationPolicy}
	 */
	public PortAllocationPolicy getPortAllocationPolicy()
	{
		return portAllocationPolicy;
	}

	/**
	 * @return the {@link PortMappingPolicy}
	 */
	public PortMappingPolicy getPortMappingPolicy()
	{
		return portMappingPolicy;
	}

	private int getRandomPort_UnusedIfPossible()
	{
		boolean found = false;
		int ret = 0;
		int limit = pToE.size();
		int counter = 0;
		do
		{
			ret = 1 + portRandom.nextInt(MAX_PORT);
			counter++;
			if (!pToE.containsKey(ret) || counter > limit)
				found = true;
		} while (!found);
		return ret;
	}

	/**
	 * @return the clashingOverrides
	 */
	public boolean isClashingOverrides()
	{
		return clashingOverrides;
	}

	public boolean isExpirationTimerEnabled()
	{
		return isExpirationTimerEnabled;
	}

	private boolean isExpired(int port)
	{
		boolean contains = timers.containsKey(port);
		return isExpirationTimerEnabled
		&& contains
		&& (System.currentTimeMillis() - timers.get(port) > EXPIRATION_TIMER);
	}

	private void logDropped(InetAddress srcIp, int srcPort, String string)
	{
		log.debug(toString() + ", Dropped=" + srcIp + ":" + srcPort + "Reason:"
				+ string);
		// log.debug("{}, Dropped= {}:{} Reason,{}", new
		// Object[] { toString(),
		// srcIp, srcPort, string });
	}

	private int mapPort(
			Endpoint v,
			InetAddress dstIp,
			int dstPort,
			InetAddress srcIp,
			int srcPort)
	{
		int portToMap = 0;
		// Check if others mapped that port
		switch (portAllocationPolicy)
		{
		case PORT_PRESERVATION:
			Endpoint mappedEnd = pToE.get(srcPort);
			if (mappedEnd == null)
				// Nobody is using that port
				portToMap = srcPort;
			else
			{
				log.trace("CL:bw=" + mappedEnd + "&" + v);
				// Somebody is already using that port
				if (clashingOverrides || isExpired(srcPort))
					/*
					 * Remove mapping from NAT table
					 */
					// Mapping will be automatically
					// overridden
					portToMap = srcPort;
				else
					switch (alternativePortAllocationPolicy)
					{
					case PORT_CONTIGUITY:
						portToMap = getContiguousPort();
						break;
					case RANDOM:
						portToMap = getRandomPort_UnusedIfPossible();
						break;
					default:
						break;
					}
			}
			break;
		case PORT_CONTIGUITY:
			portToMap = getContiguousPort();
			break;
		case RANDOM:
			portToMap = getRandomPort_UnusedIfPossible();
			break;
		default:
			break;
		}
		// Map the new one, mapping overrides current
		// mapping
		pToE.put(portToMap, v);
		Map<InetAddress, Map<Integer, Integer>> map = null;
		if (!vToDs.containsKey(v))
		{
			map = new HashMap<InetAddress, Map<Integer, Integer>>();
			vToDs.put(v, map);
		} else
			map = vToDs.get(v);
		Map<Integer, Integer> set = null;
		if (!map.containsKey(dstIp))
		{
			set = new HashMap<Integer, Integer>();
			map.put(dstIp, set);
		} else
			set = map.get(dstIp);
		set.put(dstPort, portToMap);
		return portToMap;
	}

	/**
	 * Emulates receival of packet on the NAT's external
	 * interface
	 * 
	 * @param srcIp
	 *            The source {@link InetAddress}
	 * @param srcPort
	 *            The source port
	 * @param dstIp
	 *            The destination {@link InetAddress},
	 *            specifically the IP of this instance of
	 *            the NAT box
	 * @param dstPort
	 *            The destination port
	 * 
	 * @return the {@link InetSocketAddress} of the
	 *         destination peer in the private network, null
	 *         of packet is dropped
	 */
	public InetSocketAddress receivedOnExternal(
			InetAddress srcIp,
			int srcPort,
			InetAddress dstIp,
			int dstPort)
	{
		boolean forward = false;
		// InetAddress dstIp =
		// packet.getDestinationAddress();
		// InetAddress srcIp = packet.getSourceAddress();
		// int dstPort = packet.getDestinationPort();
		// int srcPort = packet.getSourcePort();
		Endpoint d = new Endpoint();
		d.setAddr(srcIp);
		d.setPort(srcPort);
		Endpoint v = pToE.get(dstPort);
		if (v != null && !isExpired(dstPort))
		{
			// Mapping exists
			Map<InetAddress, Map<Integer, Integer>> map = vToDs.get(v);
			switch (filteringPolicy)
			{
			case ENDPOINT_INDEPENDENT:
				forward = true;
				break;
			case HOST_DEPENDENT:
				if (map.containsKey(srcIp))
				{
					Map<Integer, Integer> portMap = map.get(srcIp);
					Collection<Integer> ints = portMap.values();
					if (ints.contains(dstPort))
						forward = true;
					else
					{
						/* No port is open for that peer */
					}
				}
				break;
			case PORT_DEPENDENT:
				if (map.containsKey(srcIp))
				{
					Map<Integer, Integer> portMap = map.get(srcIp);
					Collection<Integer> ints = portMap.values();
					if (ints.contains(dstPort) && portMap.get(srcPort) != null
							&& portMap.get(srcPort) == dstPort)
						forward = true;
					else
					{
						/*
						 * Either no port is open for that
						 * peer or src port is not the same
						 * as it sent earlier
						 */
						// Drop
					}
				} else
				{
					// Drop
				}
				break;
			default:
				break;
			}
		}
		if (forward)
			// Change dest port
			return new InetSocketAddress(v.getAddr(), v.getPort());
		else
		{
			logDropped(srcIp, srcPort, "filtering");
			return null;
		}
	}

	/**
	 * Emulates receival of packet on the NAT's internal
	 * interface
	 * 
	 * @param srcIp
	 *            The source {@link InetAddress}
	 * @param srcPort
	 *            The source port
	 * @param dstIp
	 *            The destination {@link InetAddress},
	 *            specifically the IP of this instance of
	 *            the NAT box
	 * @param dstPort
	 *            The destination port
	 * 
	 * @return the {@link InetSocketAddress} of the
	 *         destination peer in the public network
	 */
	public InetSocketAddress receivedOnInternal(
			InetAddress srcIp,
			int srcPort,
			InetAddress dstIp,
			int dstPort)
	{
		// Start GC if first
		if (isExpirationTimerEnabled && pToE.isEmpty())
			startGC();
		// InetAddress dstIp =
		// packet.getDestinationAddress();
		// InetAddress srcIp = packet.getSourceAddress();
		// int dstPort = packet.getDestinationPort();
		// int srcPort = packet.getSourcePort();
		int mappedPort = srcPort;
		// Reused endpoints
		Endpoint v = new Endpoint();
		v.setAddr(srcIp);
		v.setPort(srcPort);
		// Reused endpoints
		Endpoint d = new Endpoint();
		d.setAddr(dstIp);
		d.setPort(dstPort);
		if (!vToDs.containsKey(v))
			// Not talking with the host
			/*
			 * There is no mapping for this two endpoints,
			 * we need to open a port
			 */
			mappedPort = mapPort(v, dstIp, dstPort, srcIp, srcPort);
		else
		{
			// If v is already talking with somebody outside
			Map<InetAddress, Map<Integer, Integer>> ds = vToDs.get(v);
			if (ds.containsKey(dstIp) && ds.get(dstIp).containsKey(dstPort))
				mappedPort = ds.get(dstIp).get(dstPort);
			else
				// Allocation Policy Matters
				switch (portMappingPolicy)
				{
				case ENDPOINT_INDEPENDENT:
					// To E corresponds always a single U
					mappedPort = ds.values().iterator().next().values()
					.iterator().next();
					break;
				case HOST_DEPENDENT:
					if (ds.containsKey(dstIp))
						mappedPort = ds.get(dstIp).values().iterator().next();
					else
						// Allocate a new port
						mappedPort = mapPort(v, dstIp, dstPort, srcIp, srcPort);
					break;
				case PORT_DEPENDENT:
					// Allocate a new port. Case in which
					// mapping is already present was
					// considered previously
					mappedPort = mapPort(v, dstIp, dstPort, srcIp, srcPort);
					break;
				default:
					break;
				}
		}
		// Create/Renew Timestamp
		timers.put(mappedPort, System.currentTimeMillis());
		// Forward
		return new InetSocketAddress(publicAddress, mappedPort);
	}

	/**
	 * Resets this instance's state
	 */
	public void reset()
	{
		// in_out = new HashMap<Endpoint, Set<Endpoint>>();
		pToE = new HashMap<Integer, Endpoint>();
		vToDs = new HashMap<Endpoint, Map<InetAddress, Map<Integer, Integer>>>();
		timers = new HashMap<Integer, Long>();
		// NATmappings = new HashMap<EndpointCombination,
		// Integer>();
		portRandom = new Random(RANDOM_PORT_SEED);
		portCounter = portRandom.nextInt(MAX_PORT);
	}

	/**
	 * @param clashingOverrides
	 *            the clashingOverrides to set
	 */
	public void setClashingOverrides(boolean clashingOverrides)
	{
		this.clashingOverrides = clashingOverrides;
	}

	/**
	 * Enables NAT mappings' timer expiration in this
	 * instance of the class
	 * 
	 * @param isExpirationTimerEnabled
	 */
	public void setExpirationTimerEnabled(boolean isExpirationTimerEnabled)
	{
		this.isExpirationTimerEnabled = isExpirationTimerEnabled;
	}

	public void setFilteringPolicy(FilteringPolicy policy)
	{
		filteringPolicy = policy;
	}

	public void setPortAllocationPolicy(
			PortAllocationPolicy portMappingTechnique)
	{
		portAllocationPolicy = portMappingTechnique;
	}

	public void setPortMappingPolicy(PortMappingPolicy portAllocationPolicy)
	{
		portMappingPolicy = portAllocationPolicy;
	}

	private void startGC()
	{
		gc.schedule(new TimerGarbageCollector(), TIMER_GC_PERIOD);
	}
}
